"""
作者：鑫鑫鑫
微信：xinxinjijin8
WQid：XX42289
日期：2025.01.02
未经作者允许，请勿转载
"""
from machine_lib import (async_set_alpha_properties,
                         async_login, get_alphas, prune,
                         trade_when_factory)
from config import *
import asyncio
import aiofiles
import time
from collections import defaultdict


class SessionManager:
    def __init__(self, session, start_time, expiry_time):
        self.session = session
        self.start_time = start_time
        self.expiry_time = expiry_time

    async def refresh_session(self):
        print("Session expired, logging in again...")
        await self.session.close()
        self.session = await async_login()
        self.start_time = time.time()


async def simulate_multi(session_manager, alpha_expression_list: list, region_info, name, neut, decay, delay, stone_bag,
                         tags=['None'], semaphore=None):
    """
    单次模拟一个alpha表达式对应的某个地区的信息
    """
    brain_api_url = 'https://api.worldquantbrain.com'

    async with semaphore:
        # 每个任务在执行前都检查会话时间
        if time.time() - session_manager.start_time > session_manager.expiry_time:
            await session_manager.refresh_session()

        if len(alpha_expression_list) > 10:
            raise ValueError("The number of alpha expressions in a pool should be less than 10")

        region, uni = region_info

        # 产生一个pool，一个pool里最多10个alpha
        sim_data_list = []
        for alpha_expression in alpha_expression_list:
            alpha = "%s" % (alpha_expression)
            print("Simulating for alpha: %s, region: %s, universe: %s, decay: %s, delay: %s" % (alpha, region, uni, decay, delay))

            simulation_data = {
                'type': 'REGULAR',
                'settings': {
                    'instrumentType': 'EQUITY',
                    'region': region,
                    'universe': uni,
                    'delay': delay,
                    'decay': decay,
                    'neutralization': neut,
                    'truncation': 0.08,
                    'pasteurization': 'ON',
                    'unitHandling': 'VERIFY',
                    'nanHandling': 'ON',
                    'language': 'FASTEXPR',
                    'visualization': False,
                },
                'regular': alpha
            }
            sim_data_list.append(simulation_data)

        # 一次性提交10个alpha作为单个task
        while True:
            try:
                async with session_manager.session.post('https://api.worldquantbrain.com/simulations',
                                                        json=sim_data_list) as simulation_response:
                    simulation_progress_url = simulation_response.headers.get('Location', 0)
                    if simulation_progress_url == 0:
                        json_data = await simulation_response.json()
                        if type(json_data) == list:
                            print(json_data)
                            detail = json_data.get("detail", 0)
                        else:
                            detail = json_data.get("detail", 0)
                        if detail == 'SIMULATION_LIMIT_EXCEEDED':
                            print("Limited by the number of simulations allowed per time")
                            await asyncio.sleep(5)
                        else:
                            print("detail:", detail)
                            print("json_data:", json_data)
                            print("Alpha expression is duplicated")
                            await asyncio.sleep(1)
                            return 0
                    else:
                        print('simulation_progress_url:', simulation_progress_url)
                        break
            except Exception as e:
                print(f"error occur: {e}")
                await asyncio.sleep(60)
                return

        while True:
            try:
                async with session_manager.session.get(simulation_progress_url) as resp:
                    json_data = await resp.json()
                    # 获取响应头
                    headers = resp.headers
                    retry_after = headers.get('Retry-After', 0)
                    if retry_after == 0:
                        status = json_data.get("status", 0)
                        children = json_data.get("children", [])
                        if status == 'ERROR':
                            print("Error in simulation:", simulation_progress_url)
                        elif status != "COMPLETE":
                            print("Not complete : %s" % (simulation_progress_url))
                            async with session_manager.session.delete(simulation_progress_url) as delete_resp:
                                # print(delete_resp.status_code)
                                delete_json_data = await delete_resp.json()
                                if delete_json_data.get("detail", 0) == "未找到。":
                                    print("succeed in delete:", simulation_progress_url)
                                else:
                                    print("failed in delete:", simulation_progress_url)
                        else:
                            print('simulation_progress_url completed!', simulation_progress_url)
                        break
                    await asyncio.sleep(float(retry_after))

                    # print("Sleeping for " + simulation_progress.headers["Retry-After"] + " seconds")
            except Exception as e:
                print("Error while checking progress:", str(e))
                await asyncio.sleep(60)

        # alpha_id = simulation_progress.json()["alpha"]
        children_list = []
        for child in children:
            try:
                async with session_manager.session.get(brain_api_url + "/simulations/" + child) as child_progress:
                    json_data = await child_progress.json()
                    alpha_id = json_data["alpha"]
                    alpha_express = json_data["regular"]
                    # print("这里是我想找的:", json_data)

                    await async_set_alpha_properties(session_manager.session,
                                                     alpha_id,
                                                     name="%s" % name,
                                                     color=None,
                                                     tags=tags)

                    # 将alpha保存到文件
                    async with aiofiles.open(f'records/{name}_simulated_alpha_expression.txt', mode='a') as f:
                        await f.write(alpha_express + '\n')

            except KeyError:
                print("Failed to retrieve alpha ID for: %s" % (brain_api_url + "/simulations/" + child))
            except Exception as e:
                print("An error occurred while setting alpha properties:", str(e))

        return 0



async def simulate_multiple_tasks(alpha_list, region_list, decay_list, delay_list, name, neut, stone_bag, n=10):
    semaphore = asyncio.Semaphore(n)
    tasks = []
    tags = [name]

    session_managers = []
    for _ in range(n):
        # 记录登录的开始时间，并为每个 session_manager 创建独立的 session
        session_start_time = time.time()
        session = await async_login()
        session_expiry_time = 3 * 60 * 60  # 3 小时
        session_manager = SessionManager(session, session_start_time, session_expiry_time)
        session_managers.append(session_manager)

    if region_list[0][0] == "GLB":
        alpha_list = [alpha_list[i:i + 5] for i in range(0, len(alpha_list), 5)]
    else:
        alpha_list = [alpha_list[i:i + 10] for i in range(0, len(alpha_list), 10)]

    # 将任务划分成 n 份
    chunk_size = (len(alpha_list) + n - 1) // n  # 向上取整
    task_chunks = [alpha_list[i:i + chunk_size] for i in range(0, len(alpha_list), chunk_size)]
    region_chunks = [region_list[i:i + chunk_size] for i in range(0, len(region_list), chunk_size)]
    decay_chunks = [decay_list[i:i + chunk_size] for i in range(0, len(decay_list), chunk_size)]
    delay_chunks = [delay_list[i:i + chunk_size] for i in range(0, len(delay_list), chunk_size)]

    for i, (alpha_chunks, region_chunk, decay_chunk, delay_chunk) in enumerate(
            zip(task_chunks, region_chunks, decay_chunks, delay_chunks)):
        # 获取当前 chunk 对应的 session_manager
        current_session_manager = session_managers[i]
        for alpha_chunk, region, decay, delay in zip(alpha_chunks, region_chunk, decay_chunk, delay_chunk):
            # 将任务与当前的 session_manager 关联
            task = simulate_multi(current_session_manager, alpha_chunk, region, name, neut, decay, delay, stone_bag,
                                  tags, semaphore)
            tasks.append(task)

    await asyncio.gather(*tasks)

    # 关闭所有会话
    for session_manager in session_managers:
        await session_manager.session.close()


def read_completed_alphas(filepath):
    """
    从指定文件中读取已经完成的alpha表达式
    """
    completed_alphas = set()
    try:
        with open(filepath, mode='r') as f:
            for line in f:
                completed_alphas.add(line.strip())
    except FileNotFoundError:
        print(f"File {filepath} not found.")
    return completed_alphas



if __name__ == '__main__':
    n_jobs = 1  # 多线程

    # 遍历 RECORDS_PATH 下的所有文件，获取已完成的alpha表达式
    step2_tag_list = []
    for file in os.listdir(RECORDS_PATH):
        if file.endswith("step2_simulated_alpha_expression.txt"):
            step2_tag = file.replace("_simulated_alpha_expression.txt", "")
            step2_tag_list.append(step2_tag)

    if len(step2_tag_list) == 0:
        print("No step2 tags found.")

    for step2_tag in step2_tag_list:
        if 'ILLIQUID' in step2_tag:
            region, delay, instrumentType, ILLIQUID_tag, universe, dataset_id, step_num = step2_tag.split('_')
            universe = f'{ILLIQUID_tag}_{universe}'
        else:
            region, delay, instrumentType, universe, dataset_id, step_num = step2_tag.split('_')

        if int(delay) == 1:
            sharpe_step3_th = 1.2
            fitness_step3_th = 0.75
        elif int(delay) == 0:
            sharpe_step3_th = 2.6
            fitness_step3_th = 1.4
        else:
            print("delay must be 0 or 1.")
            continue

        step3_tag = step2_tag.replace('_step2', '_step3')

        so_tracker = get_alphas("2024-10-07", "2029-12-31",
                                sharpe_step3_th, fitness_step3_th,
                                100, 100,
                                region, universe, delay, instrumentType,
                                500, "track", tag=step2_tag)

        so_layer = prune(so_tracker['next'] + so_tracker['decay'],
                         dataset_id, 3)

        th_alpha_list = []
        for expr, decay in so_layer:
            for alpha in trade_when_factory("trade_when", expr, region, delay):
                th_alpha_list.append((alpha, decay))
        print(len(th_alpha_list))

        # 读取已完成的alpha表达式
        completed_alphas = read_completed_alphas(f'records/{step3_tag}_simulated_alpha_expression.txt')

        raw_alpha_list = th_alpha_list
        # 排除已完成的alpha表达式
        alpha_list = [alpha_decay for alpha_decay in raw_alpha_list if alpha_decay[0] not in completed_alphas]
        if len(alpha_list) == 0:
            print(step3_tag, "is done.")
            continue
        print(step3_tag, "progress: {}/{}".format(len(raw_alpha_list) - len(alpha_list), len(raw_alpha_list)))

        grouped_dict = defaultdict(list)
        for alpha, decay in alpha_list:
            grouped_dict[decay].append(alpha)

        for decay in grouped_dict:
            alpha_list = grouped_dict[decay]
            decay_list = [decay] * len(alpha_list)
            delay_list = [delay] * len(alpha_list)
            region_list = [(region, universe)] * len(alpha_list)  # 扩展 region_list
            
            # USA,EUR,ASI and CHN
            if region in ['USA', 'EUR', 'ASI', 'CHN']:
                neut = 'SUBINDUSTRY'
            elif region in ['GLB']:
                neut = "SUBINDUSTRY"
            elif region in ['AMR']:
                neut = "SUBINDUSTRY"
            else:
                neut = 'SUBINDUSTRY'
            
            # 执行异步模拟，并控制并发数量为3
            asyncio.run(simulate_multiple_tasks(alpha_list, region_list, decay_list, delay_list,
                                                step3_tag, neut,
                                                [], n=n_jobs))

    print("All done. Sleep 600s...")
    time.sleep(600)
    print("Wake up.")

